
import torch
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.nn.modules.module import Module


class GNNLayer(Module):
    """
    图卷积模型
    """
    def __init__(self, in_features, out_features):
        """
        网络结构
        Args:
            in_features: 输入数据的特征
            out_features: 输出数据的特征
        """
        super(GNNLayer, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
        torch.nn.init.xavier_uniform_(self.weight)
        # torch.nn.init.xavier_normal_(self.weight)

    def forward(self, features, adj, active=True):
        """
        网络前向传递机制
        Args:
            features: 输入数据
            adj: 数据对应的邻接矩阵
            active: 是否需要激活函数
        """
        support = torch.mm(features, self.weight)
        output = torch.mm(adj, support)
        if active:
            output = F.relu(output)
        return output
